[None][feat] Self-sampling GVR V2 prefill indexer top-K - #18702

Description

Stacked on #18446. Extends the hint-free self-sampling GVR V2 indexer top-K
(introduced for decode by #18410/#18446) to the prefill phase, so a DSA
model that selects self-sampling for decode uses the same engine for prefill
instead of the CUDA radix path — one config, one kernel family, one warmup.

Dispatch is unchanged and needs no new config/API field: prefill uses the
self-sampling engine on exactly the layers where the two-level dispatch
(enable_heuristic_topk × use_self_sampling_topk) already picks it for
decode. The temporal-hint (V1) layers keep the exact radix prefill (that
engine has no prefill form). The TopK.forward(..., is_prefill=True, row_starts=..., row_ends=...) interface is unchanged.

enable_heuristic_topk=false (default)     -> exact insertion/radix prefill
enable_heuristic_topk=true
  use_self_sampling_topk=true (default)   -> self-sampling GVR prefill (this PR)
  use_self_sampling_topk=false            -> exact radix prefill (temporal has no prefill form)

Kernel (gvr_topk_decode_self_sampling.py)

A constexpr prefill mode on GvrMainKernel only (the streaming family, which
is correct for any row length). Per row r it selects the top-K of
logits[r, ks:ke] (ks=row_starts[r], ke=row_ends[r], compressed column
units), writes the LOCAL frame (column - ks) with a trailing -1 pad, and
emits identity 0..nv-1 when nv = ke-ks <= k — matching the radix short-row
contract. ks/ke ride the otherwise-unused pre_idx/kv_lens ABI slots
(hint-free never reads them), so the kernel signature is byte-identical; the
prefill compile is a distinct cache key.

The base is rounded down to a 16B boundary (col0 = ks & ~3) so the vector
loads stay aligned; the ≤3 lead = ks & 3 lanes are masked positionally (a
one-bit clear of the i0==0 classification mask, a valid-lane substitution in
the sample fold, and i >= lead guards on the two whole-row escape sweeps),
never with a materialized -inf (which would drive the sample histogram index
out of bounds and mis-rank a genuine -inf tie class). The row-extent clamps
tighten to the last in-window float4, so the kernel reads exactly
[r*npad + col0, r*npad + ke) with no dependence on producer slack. Output
indices subtract lead at the six write sites. Every edit is const_expr-gated
on the prefill flag; the decode/legacy codegen is unchanged (the full decode
suite passes bit-for-bit).

Host / module / framework

• run_prefill(logits, row_starts, row_ends, indices, max_row_len=None) — same
  validation style as run_varlen but keyed on stride(0) for all row counts
  (DeepGEMM prefill rows are 1024B-aligned with ≥256-float slack, so the
  single-row shape[1] rule that would reject odd-width tiles does not apply),
  no device reads (envelope from shape[1]), and it slabs launches to ≤32768
  rows so gridDim.y never overflows. Forces R=1 (single CTA per row, no
  SPLIT/workspace); the launcher cache is keyed on (row tier, k, envelope bucket) so arbitrary q-tile row counts do not grow it.
• warmup_prefill — compiles the ≤6 engines per k before serving (placed
  before the DeepGEMM decode-stride guard in warmup_selfsampling_topk, so an
  odd msl_c cannot skip it); best-effort under the existing OOM guard.
• TopK._forward_prefill gains a CUTE_DSL_GVR branch: all-short tiles
  (shape[1] <= k) take the exact radix identity path silently; the hardware-
  format gate (fp32, unit inner stride, float4 row stride, 16B base) routes to
  the engine with a one-time info_once; a miss falls back to radix with a
  warning_once. CUTE_DSL_RADIX still raises NotImplementedError.
• The indexer selects CUTE_DSL_GVR for prefill iff it selects it for decode
  and the self-sampling prerequisites hold; no llm_args/API change.

Prefill top-K is never CUDA-graph captured (mla_dsa_attn_inplace is a
piecewise boundary op); there are no host syncs; PDL parity with the radix
prefill is none. Disaggregated context-only servers are covered by the
unconditional warmup hook. The port is compression-ratio agnostic (ks/ke
are already compressed), so DeepSeek-V3.2 (cr=1, K=2048), V4 Flash (cr=4,
K=512) and V4 Pro (cr=4, K=1024) share one prefill path.

Validation

Kernel/host exactness against a tie-aware radix-parity reference on a B200
(SM100), per-row [ks, ke) windows on DeepGEMM-strided storage:

• causal ramps straddling k; multi-request packed windows with ks % 4 in
  {1,2,3} and +inf/NaN/3e38/-inf poison outside the window; short rows
  (nv in {0, 1, k-1, k, k+1}); all-equal / two-valued degenerate paths;
  odd-width single-row tiles; the -inf tie class crossed with a misaligned
  lead (asserts no negative index leaks); a >65535-row slab. K in
  {512, 1024, 2048}.
• The full decode self-sampling suite passes unchanged (legacy codegen intact).
• Module dispatch (CPU): self-sampling routing, format-gate and all-short
  fallthrough to radix, temporal-mode fallthrough, output-width assertion.
• Warmup: idempotent, and zero re-JIT across row counts {1..4096} × envelope
  buckets after warmup; the prefill compile key is distinct from the decode
  varlen key.
• Indexer/metadata selection: prefill engine chosen iff decode is; off-envelope
  layers keep radix; the warmup prefill leg fires (including for an odd
  msl_c).

GSM8K accuracy — end-to-end parity (B200 TEP8, 5-shot)

End-to-end trtllm-eval gsm8k (tp8 ep8, one served config per arm). An
armproof gate asserts each arm actually engaged the intended top-K path —
decode/prefill engagement is grepped from the served logs (viol=0), so the
"V2 prefill" numbers are verified to run the new kernel, not silently fall back
to radix. The self-sampling prefill (this PR) matches both the exact radix
baseline and the pre-PR "V2 decode + radix prefill" state within GSM8K sampling
noise (stderr ±~0.6). Scores are exact_match averaged over flexible-extract +
strict-match, normalized to 0–100.

All three pf_on runs are confirmed prefill_engaged=yes — the V2 prefill
kernel actually ran and the score still lands on top of both baselines, so the
prefill top-K swap is accuracy-neutral across DSV3.2, V4-Flash and V4-Pro. Each
model runs its own index_topk (Flash 512 / Pro 1024 / DSV3.2 2048).

Preliminary kernel microbenchmark

A single-B200 cold-L2 CUDA-event sweep (random Gaussian logits, 20 reps/shape) pairs
the radix topKPerRowPrefill against the DSL prefill kernel across 14 prefill shapes
(DSv3.2 / V4, K in {512,1024,2048}, rows 64..32768, n_kv 4096..262144). The
self-sampling kernel is faster on every shape (median ~4×, worst 0.859×), so no
N-gate to radix is added. This is a gross-regression sanity check on random data;
the in-pipeline nsys numbers below are what govern the ship decision.

Operator-level top-K speedup vs the radix CUDA path

Test env: B200 (SM100), 8 GPUs, tp8 ep8, PyTorch backend, chunked prefill
(chunk=8192), KV fp8. An nsys A/B — pf_on GVR V2 selection kernel GvrMainKernel
vs pf_off radix topKPerRowPrefill<K> — isolated by an exact kernel-name
set-difference (~103–110 kernel types per run; the two selection kernels are the
only differentiator, verified). Numbers are rank-0 GPU kernel duration of the
prefill window (NVTX-validated), dummy-AR excluded from the DEP8 denominators.

Per-prefill-window rank-0 kernel time (radix baseline → V2), then the speedup:

• The V2 selection kernel is 1.84–2.61× faster than the production radix kernel
  (e.g. Pro DEP8/990K: radix 3071 ms → V2 1177 ms = 2.61×). The ratio grows with
  context/candidate count (radix topKPerRowPrefill scales worse with K and candidates;
  V2 self-sampling is ~K-invariant — Pro K=512→1024 moved it only 2.68×→2.61×).

• The share is config-dependent: for the same Flash it goes 0.60%→3.43% as ISL
  grows (262K→990K) and attention-DP removes the ~33% NCCL comms from the denominator
  (TEP8→DEP8). Because top-K selection stays ≤3.4% of prefill, the 1.84–2.61× kernel
  win translates to a modest few-% end-to-end gain at throughput scale (serving results
  below), not a large e2e speedup.

• Where the prefill time goes (DEP8/990016, dummy-AR-excluded), rank-0 kernel
  work-share — the dominant kernel is the indexer MQA logits, not the top-K
  selection, and at DEP8 there is no real TP AllReduce:

  prefill layer	Flash (K512)	Pro (K1024)
  indexer MQA logits	34.0%	23.7%
  attention / MLA (fmha + mhc)	24.5%	27.0%
  dense GEMM (QKV/O) + norm/rope/quant	23.0%	29.0%
  MoE (A2A + expert GEMM + finalize + gate)	13.5%	16.8%
  indexer top-K selection (this PR)	3.4%	2.4%
  indexer other (gather / compressor / local_to_global)	1.5%	1.1%
  TP comms (non-dummy)	0%	0%

  Full per-kernel breakdowns + all trace paths:
  indexer_topk_op_bench/op51_gvr_v2_prefill/GVR_V2_TOPK_E2E_AND_OPERATOR_REPORT.html.

• No DSV3.2 (cr=1) layered breakdown: its operator share is unmeasurable under nsys
  in this environment — profiling the DSA algorithm under attention-DP deadlocks CUPTI
  (verified it runs cleanly without nsys; TEP8+nsys separately hits an NVRTC
  cuda.h gap). Its e2e numbers (below) are unaffected. By cr=1 scaling (4× candidates
  vs cr=4) its indexer-logits + top-K shares would be higher than Flash/Pro.

End-to-end serving throughput (max-config: DEP8, long-ISL, batched, 3 reps)

Because top-K selection is only 2.4–3.4% of prefill, its end-to-end effect surfaces in a
throughput config, not a BS=1 latency one (BS=1 is TTFT-neutral — full table in the
report). This config maximizes the prefill top-K share: attention-DP (DEP8, no TP
AllReduce), each model's max ISL, batched. Test env: B200 (SM100) 8 GPUs, tp8 ep8 +
enable_attention_dp=true, chunked prefill, KV fp8, OSL=512, 3-rep median, all arms
armproof-verified; concurrency = the max that fits (Flash/DSV3.2 32, Pro 8). The PR's
change is radix prefill (pf_off) → V2 prefill (pf_on) — both share V2 decode, so this
isolates the prefill top-K engine. Raw values first, then the speedup:

The V2 prefill port gives a +2.9–4.5% throughput / −3.0–4.4% TTFT end-to-end gain over
radix prefill across all three models — a real (if modest) benefit at scale, consistent
with the 2.6× kernel speedup applied to top-K's ≤3.4%-of-prefill share. (Against the exact
radix decode+prefill baseline, V2 overall is ~neutral here, 0.998–1.014× throughput — at
this high concurrency the prefill-interleaved pipeline masks the decode benefit; per-arm
tables in the report.)

Test configuration (max-config e2e)

B200 (SM100) × 8, PyTorch backend, trtllm-bench throughput --streaming, upstream-HF
FP4+FP8-mixed weights, 3 reps (median). Arms: radix (enable_heuristic_topk=false) ·
pf_off (V2 decode + radix prefill, =true + TRTLLM_GVR_PREFILL_TOPK=0, an e2e-harness-only overlay switch that is not part of this PR) · pf_on
(V2 decode + V2 prefill, =true).

The BS=1 latency e2e (SWE-bench-64K + random, ISL 262144/65536, TEP8) — where prefill
top-K is TTFT-neutral and the V2 decode benefit shows in TPOT — is in the standalone
report.

Status

Done: kernel + host + module + indexer + metadata (decode codegen unchanged);
kernel/host exactness (39 cases) + prefill UT (28) + CPU dispatch + warmup
no-re-JIT. End-to-end validated on B200 (8 GPU) across DSV3.2 / V4-Flash / V4-Pro:
GSM8K parity (accuracy-neutral); nsys operator profile (top-K selection 1.84–2.61×
the radix kernel, ≤3.4% of prefill, MQA logits dominant); BS=1 serving TTFT-neutral;
and a max-config (DEP8, long-ISL, batched) throughput run where the prefill port adds
+2.9–4.5% throughput / −3.0–4.4% TTFT.

Rebased onto main after #18446 merged (85d9871); the diff is the prefill
delta only. Follow-ups folded in after review: the warmup hook builds the DSA
metadata itself when no warmup forward ran (context-only / draft / guided-
decoder engines), an engine missed by warmup falls back to radix instead of
raising under CUDA graph capture, and GPU tests pin the DeepGEMM logits row
stride the engine's format gate relies on.

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

Dev Engineer Review

• Prefill now selects self-sampling GVR when dispatch enables it.
• Temporal-hint layers continue to use CUDA radix prefill.
• GVR prefill supports aligned windows, padding, local-frame indices, short rows, slabbed launches, and cached launchers.
• Validation, format gates, workspace handling, warmup, and CUDA-graph fallback are included.
• Decode behavior remains compatible.
• Warmup initializes DSA metadata when required.
• No configuration or test-list files changed.
• Review should confirm API consistency, row-length handling, fallback behavior, and prior-state requirements.

QA Engineer Review

• Added dispatch, warmup, and format-gate coverage in test_dsa_indexer.py.
• Added FP4 MQA format-gate coverage through test_fp4_mqa_logits_pass_selfsampling_prefill_format_gate.
• Updated routing, fallback, temporal-mode, short-input, prior-state, and CUDA-graph tests in test_top_k.py.
• Added prefill kernel coverage in test_gvr_selfsampling_topk.py, including alignment, padding, short rows, ties, odd widths, slab limits, engine-key separation, guards, warmup idempotence, and capture safety.
• Added DSA metadata warmup coverage in test_pytorch_model_engine.py.
• No tests/integration/test_lists/, test-db/, qa/, or waives.txt changes were identified.
• Coverage of the added and modified test functions in CI or manual test lists requires follow-up.
• Verdict: needs follow-up.